﻿function wimStatusPath {
    param($mountPath,$status="mounted") 
    Microsoft.PowerShell.Management\Join-Path $mountPath $status
}

function getMountInfo {
    param($mountPath)

    $p = wimStatusPath $mountPath "mounted"

    if (Test-Path $p) {
        return (Import-Csv -Path (wimStatusPath $mountPath "mounted"))
    } else {
        return $null
    }
}

function setMountInfo {
    param($mountPath, $sourcePath, $index, [switch] $dismount)

    if ($dismount.IsPresent) {
        Remove-Item (wimStatusPath $mountPath "mounted")
    } else {
        ([PSCustomObject] @{ImagePath=$sourcePath;Index=$Index}) | 
            Export-Csv -Path (wimStatusPath $mountPath "mounted") -Force |
            Out-Null
    }
}

# function getImageFile {
#     param($mountPath) 
#     Microsoft.PowerShell.Management\Resolve-Path (Microsoft.PowerShell.Management\Get-Content (wimStatusPath $mountPath "mounted") | Select-Object -First 1)
# }

function createFakeWim {
    param($ImagePath, $ImageCount, [switch] $Force)

    $image = Microsoft.PowerShell.Management\New-Item $ImagePath -ItemType File -Force:$Force

    for ($i = 1; $i -le $ImageCount; $i += 1) {
        Microsoft.PowerShell.Management\Add-Content -Path $ImagePath -Value "EXPORT $i"
    }

    return $image
}

function getImageIndexes {
    param($ImagePath) 
    
    $set = (Microsoft.PowerShell.Management\Get-Content -Path $ImagePath | 
            Microsoft.PowerShell.Core\Where-Object  {$_ -match "EXPORT (\d+)"} | 
            Microsoft.PowerShell.Core\ForEach-Object {$matches[1]})

    return $set
}

function setupMocks {

    $mountDir = [string] (Microsoft.PowerShell.Management\Resolve-Path TestDrive:\mount)

    Mock Get-WindowsImage {
        return @(getImageIndexes $ImagePath | 
                    ForEach-Object {
                        [PSCustomObject] @{
                            Index=$_;
                            Name="Mock image $_";
                            Description="Mock image $_";
                            Size=4GB
                            }
                    })
    }

    Mock Export-WindowsImage   {
        if ((getImageIndexes $SourceImagePath) -notcontains $SourceIndex) {throw "Image index $SourceIndex not found in $SourceImagePath"}
        
        $image = Microsoft.PowerShell.Management\New-Item $DestinationImagePath -ItemType File
        Microsoft.PowerShell.Management\add-content -Path $image -Value "EXPORT 1 $SourceImagePath $SourceIndex"
    } -verifiable

    Mock Mount-WindowsImage    {
        $mountInfo = getMountInfo -mountPath $Path

        if ((Microsoft.PowerShell.Management\Test-Path -Path $ImagePath) -eq $false) {throw "$ImagePath does not exist"}
        if ($mountInfo -ne $null) {throw ("Image {0} mounted in {1}" -f $mountInfo.ImagePath,$Path)}

        if ((getImageIndexes $ImagePath) -notcontains $Index) {throw "Image index $Index not found in $ImagePath"}

        setMountInfo -mountPath $Path -sourcePath $ImagePath -index $Index
        Microsoft.PowerShell.Management\Add-Content -Path $ImagePath -Value ("MOUNTED {0} {1}" -f $Index,$mount)
    } -verifiable

    Mock Add-WindowsDriver     {
        $mountInfo = getMountInfo -mountPath $Path

        if ($mountInfo -eq $null) {throw "no image mounted at $Path"}
        # if ((Microsoft.PowerShell.Management\Test-Path $Driver) -eq $false) {throw "no driver found at $Driver"}
        $wim = $mountInfo.ImagePath
        Microsoft.PowerShell.Management\Add-Content -Path $wim -Value ("ADD_DRIVER {0} {1}" -f $mountInfo.Index,$Driver)
    } -verifiable

    Mock Dismount-WindowsImage {
        $mountInfo = getMountInfo -mountPath $path

        if ($mountInfo -eq $null) {throw "no image mounted at $Path"}

        $wim = $mountInfo.ImagePath
        
        if ($Discard) {
           Microsoft.PowerShell.Management\Add-Content -Path $wim -Value ("DISCARD {0}" -f $mountInfo.Index)
        }
        elseif ($Save)    {
            Microsoft.PowerShell.Management\Add-Content -Path $wim -Value ("SAVE {0}" -f $mountInfo.Index)
        }

        setMountInfo -mountPath $path -dismount
    } -verifiable

    Mock Copy-Item -parameterFilter {$Destination -like "$mountDir\*"} -MockWith {
        throw "Please Call copyItemToWim instead of Copy-Item"
    }

    Mock copyItemToWim -MockWith {
        param (
            [parameter(Mandatory=$true)] 
            [ValidateSet("CopyItem", "CopyFile", "CopyDirectory", "CopyContents")] 
            $CopyType,
            [parameter(Mandatory=$true)] $SourcePath,
            [parameter(Mandatory=$true)] $ImageRoot,
            [parameter(Mandatory=$true)] $ImageDirectory,
            $NewName
            )

        $item = Get-Item $SourcePath
        $src = Get-Item -Path $SourcePath
        if ($src -is [array]) {throw "copyItemToWim: $SourcePath must resolve to a single item"}

        $mountInfo = getMountInfo $imageRoot
        if ($mountInfo -eq $null) {throw "copyItemToWim: $ImageRoot is not a mounted WIM"}

        switch ($CopyType) {
            "CopyFile" {
                if ($src.PSIsContainer -eq $true) {throw "copyItemToWim: $CopyType cannot accept directory ($SourcePath)"}
                $recurse = $false
                break;
            }

            {$_ -in ("CopyDirectory","CopyContents")} {
                if ($src.PSIsContainer -eq $false) {throw "CopyItemToWim: $CopyType requires source to be a container ($SourcePath)"}
                $recurse = $true
            }

            default {
                $recurse = $true 
                break;
            }
        }


        if ($PSBoundParameters.ContainsKey("NewName")) {
            $destinationName = $NewName
        } else {
            $destinationName = $src.Name
        }

        $destination = Join-Path $ImageRoot $ImageDirectory

        # Create the destination directory if it doesn't exist.  This simplifies
        # CopyDirectory vs. CopyContents
        # if ((Test-Path -PathType Container -Path $destination) -eq $false) {
        #     New-Item -Path $destination -ItemType Directory | Out-Null
        # }

        $destination = Join-Path $destination $(if ($NewName -ne $null) {$NewName} else {$item.Name})

        $txt = "{0} {1} {2} {3}" -f @{$true="COPYR";$false="COPY1"}[$recurse],
                                    $mountInfo.Index,
                                    $src.FullName,
                                    $destination

        Microsoft.PowerShell.Management\Add-Content -Path $mountInfo.ImagePath -Value $txt
    }

    Mock createShortcut -mockWith {}
}

function setupTestEnvironment {
    param ($FakeImageCount = 1)

    # setup working directories
    $script:inputDir = New-Item -ItemType Directory -Path TestDrive:\input -Force
    $script:bspDir = New-Item -ItemType Directory -Path TestDrive:\input\bsp -force
    $script:mountDir = New-Item -ItemType Directory -Path TestDrive:\mount -Force
    $script:outputDir = New-Item -ItemType Directory -Path TestDrive:\output -Force

    $script:outputWim = "$script:outputDir\testImage.wim"


    # setup BSP input file
    $dummyBsp = @"
<?xml version="1.0" encoding="utf-8"?>
<!-- hand written for unit testing purposes -->
<FeatureManifest xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xmlns:xsd="http://www.w3.org/2001/XMLSchema" xmlns="DevboardBspManifestSchema.xsd">
  <BasePackages>
    <PackageFile Path="%PKGDIR%\test1" Name="test1.inf" />
    <PackageFile Path="%PKGDIR%\test2" Name="test2.inf" />
    <PackageFile Path="%PKGDIR%\test3" Name="test3.inf" />
  </BasePackages>
  <Features>
    <OEM />
  </Features>
</FeatureManifest>
"@
    $script:bspIndex = New-Item -ItemType File -Path $script:bspDir\fake-bsp.xml -Value $dummyBsp -Force
    $script:inputWim = createFakeWim -ImagePath "$inputDir\input.wim" -ImageCount $FakeImageCount -Force

    setupMocks
}

function getDefaultParameters {
    return @{
        SourcePath    = "$script:inputWim";
          BspManifest = "$script:bspIndex";
          MountDir    = "$script:mountDir";
       Destination    = "$script:outputWim";
        }
}


function findCopyCommands($inputFile, $copyType = "COPY.") {

    $pattern = "$copyType "
    $pattern += [RegEx]::Escape($inputFile)

    Select-String -Path $script:OutputWim -Pattern ("$copyType \d+ " + [RegEx]::Escape($inputFile)) | 
        Select-Object -ExpandProperty Line
}